#!/usr/bin/env python
# (c) Yongjin Park, 2013
import sys
import os
import time
import random
import argparse
import subprocess as sp
import networkx as nx
from datetime import datetime
from math import log, exp, sqrt, ceil

def log_msg(msg):
    tstr = datetime.now().strftime("%c")
    print >> sys.stderr, "[%s] %s" % (tstr, msg)
    sys.stderr.flush()
    return

################################################################
# input data class
class InputData:

    def __init__(self):
        self.idx2label = {}
        self.label2idx = {}
        self.data_graph = None
        self.cum_idx = 0        
        pass


    '''
    read data graph and give consistent index
    '''
    def read_data(self, data_file):
        ext = data_file.strip().split('.')[-1]
        
        # build a graph
        G0 = nx.Graph()
        tmp = map(lambda x: x.strip().split('\t'), open(data_file))
        tmp = filter(lambda x: len(x)>1, tmp)
        
        # remove empty node name
        # but allow self-loop to include isolated vertex
        tmp = filter(lambda x: len(x[0])>0 and len(x[1])>0 and x[0] != x[1], tmp)
        
        G0.add_edges_from(filter(lambda x: len(x) == 2, tmp), weight=1.)
        for e, w in [(x[:2], x[2]) for x in tmp if len(x) == 3] :
            try:
                w = float(w)
                if w >= 0.: G0.add_edge(e[0], e[1], weight=w)
            except ValueError:
                log_msg('invalid weight (%s,%s) [%s]' % (e[0], e[1], w))
                sys.exit(1)
                pass
            pass

        # build/append idx2label and label2idx
        idx2lbl = self.idx2label
        lbl2idx = self.label2idx
        for v in sorted(G0.nodes()): #_iter():
            if not lbl2idx.has_key(v):
                self.cum_idx += 1
                id = self.cum_idx
                lbl2idx[v] = id
                idx2lbl[id] = v

        G = nx.Graph()
        for u, v, dat in G0.edges_iter(data=True):
            G.add_edge(lbl2idx[u], lbl2idx[v], weight=dat['weight'])

        # print >> sys.stderr, ' V = %d E = %d' % (G.number_of_nodes(), G.number_of_edges())

        self.data_graph = G
        return G
    
    # private helper functions
    @staticmethod
    def __get_edge_str(Gsub):
        edge_str = lambda e : '%d\t%d\t%f\n%d\t%d\t%f' % (e[0], e[1], e[2]['weight'], e[1], e[0], e[2]['weight'])
        return '\n'.join(map(edge_str, Gsub.edges_iter(data=True))) + '\n'

    @staticmethod
    def __get_loop_str(Gsub):
        vset = [v for v,d in Gsub.degree_iter() if d < 1]
        return '\n'.join(map(lambda v: '%d\t%d\t0.'%(v,v), vset)) + '\n'

    '''
    generate sub-networks of subset
    write networks' weighted pairs under out_dir
    return generated network files
    '''
    def gen_sub_data(self, out_dir, nodes=None):
        edge_str = InputData.__get_edge_str
        loop_str = InputData.__get_loop_str

        fhdr = out_dir.rstrip('/') + '/'
        fname = fhdr + 'G.wpairs'

        if nodes == None :
            H = self.data_graph
        else:
            H = self.data_graph.subgraph( nodes )

        if H.number_of_edges() < H.number_of_nodes() :
            return (fname, False)

        fh = open(fname, 'w')
        fh.write(edge_str(H))    # edge string
        fh.write(loop_str(H))    # isolated vertices
        fh.close()            
        
        return (fname, True)

    pass # End of InputData


'''
tree class to represent the search history
'''
class Tree:

    def __init__(self, depth):
        self.treeobj = nx.DiGraph()
        self.pa = {}
        self.root_idx = 0
        self.treeobj.add_node(self.root_idx)
        self.node = self.treeobj.node
        self.treeobj.node[0]['isleaf'] = True
        self.depthBase = depth
        pass

    def add_subtree_with_data(self, curr, nodeData, edgeData, group):

        T = self.treeobj

        assert( T.node[curr].get('isleaf',False) )

        # index map : local subtree idx -> global tree idx
        idxmap = {}
        def fill_idx(r, k, idxmap) :
            if r not in edgeData: return    # no data

            idxmap[r] = k

            if edgeData[r][-1] > -1: return # leaf node

            fill_idx( 2*r+1, 2*k+1, idxmap )
            fill_idx( 2*r+2, 2*k+2, idxmap )
            return

        fill_idx(0, curr, idxmap)

        group_ret = {}

        for u in edgeData:

            if edgeData[u] < 3 : continue

            left = idxmap.get( edgeData[u][0], -1 )
            right = idxmap.get( edgeData[u][1], -1 )
            k = edgeData[u][2]
            node_data = nodeData.get( u, {} )

            uCurr = idxmap.get( u, -1 )
            T.node[ uCurr ] = node_data

            # internal node
            if left > 0 and right > 0 :
                T.add_edge( uCurr, left )
                T.add_edge( uCurr, right )
                self.pa[left] = uCurr
                self.pa[right] = uCurr

            else: # leaf node                
                T.node[ uCurr ]['isleaf'] = True
                T.node[ uCurr ]['subset'] = group.get(k,[])
                group_ret[ uCurr ] = group.get(k,[])


        return group_ret

    # add subtree underneath 'curr'
    # group must take 0-based index
    def add_subtree(self, curr, group):
        treeobj = self.treeobj
        bottom = self.__add_subtree(curr, 0)

        if 'subset' in treeobj.node[curr] :
            del treeobj.node[curr]['subset']
            
        group_ret = {}
        for k, idx in enumerate(bottom):
            if k in group:
                treeobj.node[idx]['subset'] = group[k]
                group_ret[idx] = group[k]


        return group_ret

    # helper function
    # d = current depth
    def __add_subtree(self, curr, d):
        treeobj = self.treeobj
        if d == self.depthBase:
            treeobj.node[curr]['isleaf'] = True
            return [curr]
        (left, right) = (2 * curr + 1, 2 * curr + 2)
        treeobj.add_edge(curr, left)
        treeobj.add_edge(curr, right)
        self.pa[left] = curr
        self.pa[right] = curr
        treeobj.node[curr]['isleaf'] = False

        lset = self.__add_subtree(left, d + 1)
        rset = self.__add_subtree(right, d + 1)
        return lset + rset

    # pruning out unnecessary subtree(s)
    # at leaf level
    # (i) if empty -> remove -> return 0
    # at internal level
    # (iia) if no child -> remove node -> return 0
    # (iib) if one child -> transfer node -> return 1
    # (iic) if two or more -> do nothing -> return 1
    def prune_subtree(self, curr=None):
        if curr == None: curr = self.root_idx
        treeobj = self.treeobj
        pa = self.pa

        nodeobj = treeobj.node[curr]

        # at leaf level
        if nodeobj.get('isleaf',False) :
            if len(nodeobj.get('subset', {})) == 0:
                treeobj.remove_node(curr)
                del pa[curr]
                return 0
            return 1
        
        # at internal level
        valid = []
        nret = 0
        for x in treeobj.neighbors(curr):
            ret = self.prune_subtree(x)
            if ret > 0 : valid.append(x)
            nret += ret

        if nret == 0: 
            if curr != self.root_idx:
                treeobj.remove_node(curr)
                del pa[curr]
            return 0

        elif nret == 1 and curr != self.root_idx :
            # transfer 
            # copy child's edges and child's data
            # delete child node
            if curr != self.root_idx:
                pa_node = pa[curr]
            else:
                pa_node = curr
            ch_node = valid[0]

            for x in treeobj.neighbors(ch_node):
                treeobj.add_edge(curr, x)
                pa[x] = curr

            # needs deepcopy?
            treeobj.node[curr] = treeobj.node[ch_node]

            treeobj.remove_node(ch_node)
            del pa[ch_node]

            return 1
        
        return 1
    
    # just output edges
    def edges_str(self):
        ret = '\n'.join( [ '%s\t%s'%(u,v) for (u,v) in self.treeobj.edges_iter()] )
        return ret

    pass # End of class Tree




################################################################
# I / O related stuff
def group_output( input_data, btree, outFileName ):
    grp_fh = open(outFileName+'.txt', 'w')
    grp_fh.write('group_id\tmembers\n')
    grpdict_fh = open(outFileName+'.dict', 'w')
    for gg, dat in btree.node.iteritems():        
        if dat.get('isleaf', False):
            grp_arr = map(lambda x: input_data.idx2label[x], dat.get('subset', []))
            if len(grp_arr) > 0 :
                grp_fh.write(str(gg) + '\t' + ','.join(grp_arr) + '\n')
                grpdict_fh.write('\n'.join(map(lambda u: u + '\t' + str(gg), grp_arr)) + '\n')

    grp_fh.close()
    grpdict_fh.close()
    return 1

def latent_output( input_data, latent, outFileName ):
    fh = open(outFileName+'latent.txt', 'w')
    for k,pvec in latent.iteritems() :
        tmp = map(lambda (v,p): input_data.idx2label.get(v,'') + '\t' + str(k) + '\t' + str(p), 
                  pvec.iteritems())
        print >> fh, '\n'.join(tmp)
        
    fh.close()
    

def tree_output( input_data, btree, outFileName ) :
    for gg, dat in btree.node.iteritems():        
        if dat.get('isleaf', False):
            grp_arr = map(lambda x: input_data.idx2label[x], dat.get('subset', []))
            dat['subset'] = grp_arr

    nx.write_gpickle( btree.treeobj, outFileName )

    return 1

# def tree_model_output( engine, outFileName ):
#     depth = engine.TreeDepth
#     nData = engine.TreeNodeData
#     fh = open(outFileName, 'w')
    
#     print >> fh, depth

#     node_data_str = lambda nd: '\t'.join( map(str, [ nd.get('a',0.), nd.get('b',0.), 
#                                                      nd.get('edges',0.), nd.get('total',0.) ] ) )
#     # print engine.TreeNodes

#     fh.write( '\n'.join( [ node_data_str(nData.get(u,{})) for u in engine.TreeNodes ] ) + '\n' )
                           
#     fh.close()
#     pass


def tree_edge_output( btree, outFileName ) :
    tree_fh = open( outFileName, 'w')
    tree_fh.write( btree.edges_str() )
    tree_fh.close()
    return 1

def get_bottom_clusters( btree ):
    ret = {}
    for gg, dat in btree.node.iteritems():        
        if dat.get('isleaf', False):
            ret[gg] = dat.get('subset', [])
    return ret

def get_bottom_dict( btree ):

    T = btree.treeobj

    def _f(r) :
        neigh = T.neighbors(r)
        if len(neigh) == 0 :
            return [ T.node[r].get('subset',[]) ]

        tmp = [ _f(x) for x in neigh ]
        return reduce(lambda a, x: a+x, tmp, [])

    botList = _f( btree.root_idx )
    get_pair = lambda uList, k: [(u,k) for u in uList]
    ret = dict( reduce(lambda a,k: a + get_pair(botList[k], k), 
                       xrange(len(botList)), []) )
    return ret

def get_tree_depth( btree ):
    T = btree.treeobj
    def _f( r ):
        neigh = T.neighbors(r)
        if len(neigh) == 0 : return 0
        ret = max( [1+_f(x) for x in neigh] )
        return ret
    
    return _f( btree.root_idx )


################################################################
# class Engine
class Engine:

    def __init__(self, inData, args):
        
        self.EXEDIR = os.path.dirname(os.path.realpath(__file__))

        try:
            self.wd = sp.check_output(['mktemp', '-d', './tmp.hsblock.XXXXXX']).strip()
        except sp.CalledProcessError:
            log_msg('cannot create working directory')
            sys.exit(1)
            pass

        self.elapsed_time = 0.
        self.inData = inData

        binFile = self.EXEDIR + '/hsblock'
        model = 'sb' if args.no_degree_correct else 'dsb'
        alg = 'mf' if args.meanfield else 'lcvi'        
        binFile = '_'.join( [ binFile, model, alg ] )
        
        log_msg('model = ' + model + ', alg = ' + alg)
         
        self.cmd_exe = [ binFile ]                         
        self.cmd_exe += [ '-i', str(args.iter), '-e', str(args.eta) ]
        if args.full : self.cmd_exe += [ '--full' ]
        
        self.args = args
        self.argmax = None
        self.latentFile = None
        self.latent = None
        self.BayesF = None
        self.TreeDepth = None
        self.TreeFile = None
        self.TreeNodeData = None
        self.TreeEdgeData = None
        pass # end of initialization

    def __del__(self):
        try:
            sp.check_call(['rm', '-rf', self.wd])            
        except sp.CalledProcessError:
            log_msg('cannot clear out working direcotry: '+self.wd)
            sys.exit(1)
        
    def clear_wd(self):
        try:
            sp.check_call(['rm', '-rf', self.wd + '/*'])
        except sp.CalledProcessError:
            log_msg('cannot clear out working direcotry: '+self.wd)
            sys.exit(1)

    def initial_latent( self, membership ):
        ret = self.wd + '/init_latent.txt'       
        fh = open( ret, 'w' )
        fh.writelines( ['%d\t%d\t1.\n'%(u,k) for u,k in membership.iteritems()] )
        fh.close()
        return ret

    def run(self, depth, nodes=None, membership=None, repeats=True, determ=False):

        depth_adjust = depth
        if nodes != None :
            depth_adjust = int(min(depth, ceil(log(len(nodes)/self.args.ssz+1.01, 2.))))

        dataFile, validData = self.inData.gen_sub_data(self.wd, nodes)

        if not validData : 
            self.argmax = {}
            self.TreeNodes = []
            self.TreeNodeData = {}
            self.TreeEdgeData = {}
            self.BayesF = 0.
            if nodes != None: self.argmax = { 0: nodes }
            return
        
        self.clear_wd()

        result_hdr = self.wd + '/out'
        curr_cmd = []
        curr_cmd += self.cmd_exe
        if membership != None: curr_cmd += ['-l', self.initial_latent(membership)]
        curr_cmd += ['-d', str(depth_adjust), '-o', result_hdr, dataFile]

        ## deterministic update
        if self.args.determ or determ : curr_cmd += [ '--determ' ]

        if repeats    : curr_cmd += ['-r', str(self.args.rep)]

        if self.args.verbose >= 2 :
            log_msg('running variational inference with depth = ' + str(depth))

        # print curr_cmd
        errFH = sys.stderr if self.args.verbose >=2 else open(os.devnull,"w")
        outFH = sys.stdout if self.args.verbose >=2 else open(os.devnull,"w")

        try:
            # clean up before running
            sp.check_call(['rm', '-rf', self.wd+'/*'])
            # run and wait; should redirect output?            
            p = sp.Popen(curr_cmd, stdout=outFH, stderr=errFH)
            p.wait()
            pass
        except (sp.CalledProcessError, OSError, IOError) as e:
            log_msg('cannot execute command:\n'+' '.join(curr_cmd))
            sys.exit(1)
            pass
                
        # 1. argmax membership
        argmaxfile = result_hdr + '.argmax'
        argmax = {}
        for x, g in map(lambda x: x.strip().split('\t'), open(argmaxfile)):
            k = int(g)
            v = int(x)
            if not argmax.has_key(k): argmax[k] = []
            argmax[k].append(v)

        self.argmax = argmax
        
        # 2. Bayes factor
        infofile = result_hdr + '.info'
        tmp = dict(map(lambda x: x.strip().split('\t'), open(infofile)))

        self.BayesF = float(tmp['bayes_factor'])
        
        
        # 3. latent probability file
        latentfile = result_hdr + '.latent'
        self.latentFile = latentfile

        latent = {}
        for x, g, p in map(lambda x: x.strip().split('\t'), open(latentfile)) :
            k = int(g)
            v = int(x)
            r = float(p)
            if not latent.has_key(k) : latent[k] = {}
            latent[k][v] = r

        self.latent = latent

        # 4. read tree model
        treefile = result_hdr + '.tree'
        self.TreeFile = treefile
        treeLines = map(lambda x: x.strip().split('\t'), open( treefile ))

        depth = int( treeLines[0][0] )
        num_tree_nodes = int(2 ** (depth+1)) - 1
        self.TreeDepth = depth

        if len(treeLines) < 1 + 2*num_tree_nodes :
            log_msg('insufficient/invalid tree output')
            sys.exit(1)
            pass

        # edge data : idx left right leaf
        edges = map(lambda x: map(int, x), treeLines[(num_tree_nodes+1):])

        # print edges

        self.TreeNodes = map(lambda x: x[0], edges)
        self.TreeEdgeData = dict( map(lambda x: (x[0],x[1:]), edges) )

        # node data : prior#1 prior#2 stat#1 stat#2
        keys = ['a','b','edges','total']
        self.TreeNodeData = dict( [ ( edges[i][0], dict( zip(keys, map(float,x)) ) )
                                for i,x in enumerate(treeLines[1:(num_tree_nodes+1)])
                                if len(x)>3 ] )

        return # end of method


    pass # end of inference engine


################################################################
# 1. initialize by iterative bisection
# 2. determine depth and determine initial Z
# 3. fit a full model

## resolve initial latent assignment and depth
## by iterative k-section
def resolve_init_model( args, engine, nodes=None ):

    btreeIterative = Tree( depth = args.d )
    engine.run( depth = args.d, nodes=nodes, repeats=True, determ=True )
    group = btreeIterative.add_subtree_with_data( btreeIterative.root_idx, 
                                                  engine.TreeNodeData, 
                                                  engine.TreeEdgeData, 
                                                  engine.argmax )

    while len(group) > 0:
        group_next = {}

        for k in group :
            if len(group[k]) < args.ssz : continue
            engine.run( depth = args.d, nodes = group[k], repeats=True, determ=True )
            if engine.BayesF <= 1e-5 : continue
            if len(engine.argmax) <= 1 : continue

            group_sub = btreeIterative.add_subtree_with_data(k, 
                                                             engine.TreeNodeData, 
                                                             engine.TreeEdgeData, 
                                                             engine.argmax)

            for s in group_sub:
                if len(group_sub[s]) > 1:
                    group_next[s] = group_sub[s]

        group = group_next
        pass

    depth = get_tree_depth( btreeIterative )
    groupIterative = get_bottom_dict( btreeIterative )

    return (groupIterative, depth, btreeIterative)

################################################################
class InfoFile:
    def __init__(self,out_dir):
        self.out_dir = out_dir
        self.out_file = out_dir.rstrip('/') + '/' + 'info.txt'
        self.content = {}

    def __del__(self):
        fh = open(self.out_file, 'w')
        fh.writelines( [ '%s\t%s\n'%( str(k), str(self.content[k]) )
                        for k in sorted(self.content.keys()) ] )
        fh.close()

    def add_info(self, key, value):
        key = key.replace('\t',' ')
        value = value.replace('\t',' ')
        self.content[key] = value

    pass


################################################################
if __name__ == '__main__' :

    parser = argparse.ArgumentParser(description='Resolve Hierarchical Stochastic Block Models',
                                     epilog='(c) 2008-2013, Yongjin Park and Joel S. Bader, {ypark28, joel.bader}@jhu.edu')
    parser.add_argument( 'dataFile', help='A file that contains a list of pairs' )
    parser.add_argument( '-d', type=int, default=1, help='The depth of base learner (default: %(default)s)' )
    parser.add_argument( '-o', '--out', default='./out', help='Output directory (default: %(default)s)' )
    parser.add_argument( '-i', '--iter', type=int, default=100, help='Number of update iterations (default: %(default)s)' )
    parser.add_argument( '-r', '--rep', type=int, default=1, help='Number of repetition (default: %(default)s)' )
    parser.add_argument( '-e', '--eta', type=float, default=1., help='Learning rate (default: %(default)s)' )
    parser.add_argument( '-z', '--ssz', type=int, default=10, help='Stop size for iterative bisection (default: %(default)s)' )
    parser.add_argument( '--full', action='store_true', help='Full exhaustive calculation in latent updates' )
    parser.add_argument( '--determ', action='store_true', help='Deterministic latent update' )
    parser.add_argument( '--meanfield', action='store_true', help='Inference by Mean-field approximation (default: locally collapsed variational)' )
    parser.add_argument( '--no-degree-correct', action='store_true', help='Just regular stochastic block model' )
    parser.add_argument( '--level', type=int, default=1, help='Level of the tree model (default: %(default)s)' )
    parser.add_argument( '-v', '--verbose', action='count', default=0, help='(default: %(default)s)' )
    
    args = parser.parse_args()

    
    # read input data and start engine
    inData = InputData()
    log_msg('reading data: ' + args.dataFile)
    inData.read_data( args.dataFile )
    level = args.level

    engine = Engine( inData, args )

    ## output direcotry
    output_dir = args.out.rstrip('/') + '/'
    try:
        sp.check_call(['mkdir', '-p', output_dir])
    except sp.CalledProcessError:
        log_msg('cannot create output directory ' + output_dir)
        sys.exit(1)

    ## output info
    infoFile = InfoFile( output_dir )
    degFlagLog = 'sb' if args.no_degree_correct else 'dsb'
    varFlagLog = 'mean-field' if args.meanfield else 'locally collapsed'
    infoFile.add_info( 'Data Network', args.dataFile )
    infoFile.add_info( 'Degree Correction', degFlagLog )
    infoFile.add_info( 'Variational Inference', varFlagLog )
    infoFile.add_info( 'Inference Started', datetime.now().strftime("%c") )

    ################################################################
    out_hdr = output_dir
    if args.level > 1 :
        log_msg('level = 1')
        out_hdr += 'L1'

    ## initial iterative bisection
    log_msg('initial fitting by iterative bisection')
    (groupIterative, depth, btreeIterative) = resolve_init_model( args, engine )
    group_output( inData, btreeIterative, out_hdr+'iterative' )
    tree_output( inData, btreeIterative, out_hdr+'treeIterative.pickle' )
    
    ## run with full depth
    log_msg('fitting with full depth = %d'%depth)
    engine.run( depth=depth, nodes=None, membership=groupIterative, repeats=False )
    btreeFinal = Tree( depth=depth )

    group = btreeFinal.add_subtree_with_data( btreeFinal.root_idx, 
                                              engine.TreeNodeData, 
                                              engine.TreeEdgeData, 
                                              engine.argmax) 
    
    if args.level > 1 : 
        infoFile.add_info( 'Level', args.level )
        group_output( inData, btreeFinal, out_hdr )
    
    ## go deeper levels in the same way
    ## iterative then full model
    for l in xrange(1,args.level) :
        log_msg('go deeper level =' + str(l+1))
        group_next = {}
        out_hdr = output_dir + 'L%d'%(l+1)

        for k in group:
            if len(group[k]) < args.ssz : continue

            (groupIterative, depth, btreeIterative) = resolve_init_model( args, engine, nodes=group[k] )
            engine.run( depth = depth, nodes = group[k], membership=groupIterative, repeats=False )
            if args.verbose > 0 : log_msg('fitting on subnetwork with full depth = %d'%depth)

            if engine.BayesF < 0. : continue

            if len(engine.argmax) > 1 :

                group_sub = btreeFinal.add_subtree_with_data( k,
                                                              engine.TreeNodeData, 
                                                              engine.TreeEdgeData, 
                                                              engine.argmax)

                for s in group_sub:
                    if len(group_sub[s]) > 1:
                        group_next[s] = group_sub[s]

        group = group_next

        if (l+1) < args.level : group_output( inData, btreeFinal, out_hdr )
        pass

    if args.verbose > 0 : log_msg('pruning the full tree')
    btreeFinal.prune_subtree()

    infoFile.add_info( 'Inference ended', datetime.now().strftime("%c") )

    ## bottom-level group output
    group_output( inData, btreeFinal, output_dir+'final' )

    ## output tree structure
    ## note: this also converts 'subset' index to names
    tree_output( inData, btreeFinal, output_dir+'treeFinal.pickle' )

    ## tree edges
    tree_edge_output(btreeFinal, output_dir+'/treeFinal.edges')

    log_msg('Done')
    sys.exit(0)
    # end of main
